#include <iostream>
#include <string>
#include <fstream>
#include <cstdlib>
#include <chrono>
#include <map>
#include <vector>
#include "cmdline.h"
#include "ldfdetector.h"

using namespace std;

// Determine whether a flow is a long duration flow or not
inline bool isLongDurationFlow(const FlowValue& flow_value, const __time_t& LDF_threshold) {
	return (flow_value.end_time.tv_sec - flow_value.start_time.tv_sec >= LDF_threshold) ? true : false;
}

// Output Long Duration Flow
inline void OutputLDF(const FlowKey& flow_key, const FlowValue& flow_value, std::ostream& out) {
	out << flow_key << " " << flow_value << std::endl;
} 

// Cover IP address from bytes to int
uint32_t IPv4BytesToInt(const uint8_t* bytes){
    uint32_t a = bytes[3] & 0xFF;   
    a |= ((bytes[2] << 8) & 0xFF00);  
    a|= ((bytes[1] << 16) & 0xFF0000);  
    a |= ((bytes[0] << 24) & 0xFF000000);
    
    return a;
};

// Report LDF and delete not LDF
void UpdateFlowMap(std::unordered_map<FlowKey, FlowValue, FlowKeyHash>& m_flow, timespec cur_time, Options& opts, std::ostream& os) {
	decltype(cur_time.tv_sec) time_interval;
	auto iter = m_flow.begin();

	while (iter != m_flow.end()) {
		// Calculate the time interval between current time and the latest packet's arrival time
		time_interval = cur_time.tv_sec - iter->second.end_time.tv_sec;

		// If the time interval is less than the max interval time, than skip this flow
		if (time_interval < opts.max_interval) {
			iter++;
			continue;
		}

		// else, determine whether the flow is long duration flow or not
		else {
			if (isLongDurationFlow(iter->second, opts.LDF_threshold)) {
				
				OutputLDF(iter->first, iter->second, std::cout);
				OutputLDF(iter->first, iter->second, os);
			}
			m_flow.erase(iter++);
		}
	}
}

void LastCheck(const std::unordered_map<FlowKey, FlowValue, FlowKeyHash>& m_flow, const Options& opts, std::ostream& os) {
	for (auto& iter: m_flow) {
		if (isLongDurationFlow(iter.second, opts.LDF_threshold)) {
			OutputLDF(iter.first, iter.second, std::cout);
			OutputLDF(iter.first, iter.second, os);			
		}
	}
}

// Determine whether the checking time node is reached and update next check time
bool TimeToCheck(const timespec& cur_time, timespec& check_time, const __time_t& interval) {
	// Time accuracy: second. Ignore nanosecond level infomation
	if ((cur_time.tv_sec - check_time.tv_sec) >= interval){
		check_time.tv_sec += interval;
		return true;
	}
		
	return false;
}


// Get the first packet's start time
timespec GetStartTime(const char* file_path) {
    pcpp::IFileReaderDevice* reader = pcpp::IFileReaderDevice::getReader(file_path);
    timespec start_time;
    pcpp::RawPacket rawPacket;

    // verify that a reader interface was indeed created
    if (reader == NULL) {
        printf("Cannot determine reader for file type\n");
        exit(1);     
    }

    // open the reader for reading
    if (!reader->open()) {
        printf("Cannot open input data file for reading\n");
        exit(1);
    }

    reader->getNextPacket(rawPacket);
    start_time = rawPacket.getPacketTimeStamp();
    reader->close();

    // free reader memory because it was created by pcpp::IFileReaderDevice::getReader()
    delete reader;

    return start_time;
}

// Show keys and values form map
void ShowMap(std::unordered_map<FlowKey, FlowValue, FlowKeyHash>& m_flow) {
	for (auto& iter : m_flow) {
		cout << iter.first << " " << iter.second << endl;
	}
}

// Insert flow
void InsertFlow(const PacketInfo& packet_info, std::unordered_map<FlowKey, FlowValue, FlowKeyHash>& m_flow, const FlowKey& flow_key) {
	FlowValue flow_value;

	flow_value.start_time = packet_info.ts;
	flow_value.end_time = packet_info.ts;
	flow_value.packet_num = 1;
	m_flow.insert(make_pair(flow_key, flow_value));
}

// Update FlowValue
void UpdateFlowValue(const PacketInfo& packet_info, std::unordered_map<FlowKey, FlowValue, FlowKeyHash>& m_flow, const FlowKey& flow_key) {
	m_flow[flow_key].end_time = packet_info.ts;
	m_flow[flow_key].packet_num += 1;
}

// GenFlowKey PacketInfo to FlowKey
FlowKey GenFlowKey(const PacketInfo& packet_info) {
	FlowKey flow_key;

	flow_key.src_ip   = packet_info.src_ip;
	flow_key.src_port = packet_info.src_port;
	flow_key.dst_ip   = packet_info.dst_ip;
	flow_key.dst_port = packet_info.dst_port;

	return flow_key;
}

// Determine whether the flow exists in the map
inline bool exist(const FlowKey& flow_key, const std::unordered_map<FlowKey, FlowValue, FlowKeyHash>& m_flow) {
	return (m_flow.find(flow_key) != m_flow.end()) ? true : false;
}

inline void ShowPacketInfo(PacketInfo flow) {
	cout << "src IP   : " << flow.src_ip << endl
		 << "src port : " << flow.src_port << endl
		 << "dst IP   : " << flow.dst_ip << endl
		 << "dst port : " << flow.dst_port << endl 
		 << "timestamp: " << flow.ts.tv_sec << "." << flow.ts.tv_nsec << endl << endl;
}

// print all requested packets in a pcap file
int PrintPcapPackets(pcpp::PcapFileReaderDevice* reader, std::ostream* out, int packetCount) {
	// read packets from the file until end-of-file or until reached user requested packet count
	int packetCountSoFar = 0;
	pcpp::RawPacket rawPacket;

	while (reader->getNextPacket(rawPacket) && packetCountSoFar != packetCount)
	{
		// parse the raw packet into a parsed packet
		pcpp::Packet parsedPacket(&rawPacket);

		// print packet to string
		(*out) << parsedPacket.toString() << std::endl;

		packetCountSoFar++;
	}
	
	// return the nubmer of packets that were printed
	return packetCountSoFar;
}

// Get packets' information, including src_IP, src_port, dst_IP, dst_port, protocol and timestamp
int PcapPacketsHandler(pcpp::PcapFileReaderDevice* reader, Options& opts){
	pcpp::RawPacket rawPacket;
	PacketInfo packet_info;
	std::unordered_map<FlowKey, FlowValue, FlowKeyHash> m_flow;
	FlowKey flow_key;
	timespec check_time = GetStartTime(opts.input_file.c_str());
	ofstream of;
	of.open(opts.output_file.c_str());

	while (reader->getNextPacket(rawPacket)){

		if (!PacketParser(rawPacket, packet_info)) {
			// ShowPacketInfo(packet_info);
			flow_key = GenFlowKey(packet_info);

			// If the packet belongs to a existed flow, update corresponding flow's record
			if (exist(flow_key, m_flow)) {
				// update
				UpdateFlowValue(packet_info, m_flow, flow_key);
			}
			// If the packet doesn't belong to any flow, creat a new flow
			else {
				// insert
				InsertFlow(packet_info, m_flow, flow_key);
			}
			
			// If the packet's timestamp is equal to check time, update map
			if (TimeToCheck(packet_info.ts, check_time, opts.check_interval)) {
				UpdateFlowMap(m_flow, packet_info.ts, opts, of);
				if (!opts.slience)
					cout << check_time.tv_sec << endl;
			}
		}
	}

	// Update the flow map for the last time
	LastCheck(m_flow, opts, of);

	of.close();

	// ShowMap(m_flow);

	return 0;
}

// Get src_IP and dst_IP from IP layer
int IPv4LayerHandler(const pcpp::Packet parsedPacket, PacketInfo& flow) {
	pcpp::IPv4Layer* ipLayer = parsedPacket.getLayerOfType<pcpp::IPv4Layer>();

	if (ipLayer == NULL){
		// printf("Something went wrong, couldn't find IPv4 layer\n");
		return 1;
	}

	flow.src_ip = IPv4BytesToInt(ipLayer->getSrcIpAddress().toBytes());
	flow.dst_ip = IPv4BytesToInt(ipLayer->getDstIpAddress().toBytes());

	return 0;
}

// Get src_port and dst_port from TCP layer
int TCPLayerHandler(const pcpp::Packet parsedPacket, PacketInfo& flow) {
	pcpp::TcpLayer* tcpLayer = parsedPacket.getLayerOfType<pcpp::TcpLayer>();

	if (tcpLayer == NULL){
		// printf("Something went wrong, couldn't find TCP layer\n");
		return 1;
	}

	flow.src_port = ntohs(tcpLayer->getTcpHeader()->portSrc);
	flow.dst_port = ntohs(tcpLayer->getTcpHeader()->portDst);

	return 0;
}

// Get src_IP, src_port, dst_IP, dst_port, protocol, timestemp from a raw packet
int PacketParser(pcpp::RawPacket& rawPacket, PacketInfo& flow) {
    flow.ts = rawPacket.getPacketTimeStamp(); // get timestamp
    pcpp::Packet parsedPacket(&rawPacket); // parse the raw packet into a parsed packet
	int ipv4_status {0};
	int tcp_status {0};

	// bool is_bottom_ip_layer = true;
	
	// go over the layers one by one and find out its type, 
	for (pcpp::Layer* curLayer = parsedPacket.getFirstLayer(); curLayer != NULL; curLayer = curLayer->getNextLayer()) {
		pcpp::ProtocolType protocolType = curLayer->getProtocol();

		// if (protocolType == pcpp::IPv4 && is_bottom_ip_layer) {
		// 	is_bottom_ip_layer = false;
		// 	ipv4_status = IPv4LayerHandler(parsedPacket, flow);
		// }
		// if (protocolType == pcpp::TCP) {
		// 	tcp_status = TCPLayerHandler(parsedPacket, flow);
		// 	break;
		// }
		switch (protocolType){
			case pcpp::IPv4:
				ipv4_status = IPv4LayerHandler(parsedPacket, flow);
				break;
			case pcpp::TCP:
				tcp_status = TCPLayerHandler(parsedPacket, flow);
				break;
		}
	}

	return (ipv4_status || tcp_status) ? 1 : 0;
}

int main(int argc, char* argv[]) {
	cmdline::parser arg_parser;

	arg_parser.add<string>(
		"input_file",            // long name
		'i',                     // short name
		"input pcap file path", // description
		false,                   // wether this arg is mandatory
		"./input.pcap"           // default value
	);

	arg_parser.add<string>(
		"output_file",
		'o',
		"output file path. This file contains Long Duration Flows' information",
		false,
		"./output.txt"
	);

	arg_parser.add<__time_t>(
		"LDF_threshold",
		't',
		"the minimum duration of a long duration flow (second)",
		false,
		300
	);

	arg_parser.add<__time_t>(
		"max_interval",
		'x',
		"the maximum interval time (in second) between two consecutive packets of a same network traffic flow",
		false,
		30
	);

	arg_parser.add<__time_t>(
		"check_interval",
		'c',
		"the frequency to update the flow map. This argument is supposed to less than \"max_interval\"",
		false,
		1
	);

	arg_parser.add(
		"slience",
		's',
		"Avoid verbose."
	);

	arg_parser.add<int>(
		"print",
		'p',
		"print first n-th packets' information",
		false,
		0
	);

	arg_parser.parse_check(argc, argv);

	Options opts {
		arg_parser.get<string>("input_file"),
		arg_parser.get<string>("output_file"),
		arg_parser.get<__time_t>("LDF_threshold"),
		arg_parser.get<__time_t>("max_interval"),
		arg_parser.get<__time_t>("check_interval"),
		arg_parser.exist("slience")
	};

	pcpp::IFileReaderDevice* reader = pcpp::IFileReaderDevice::getReader(opts.input_file.c_str()); // open a pcap file for reading

    // verify that a reader interface was indeed created
    if (reader == NULL) {
        printf("Cannot determine reader for file type\n");
        exit(1);     
    }

    // open the reader for reading
    if (!reader->open()) {
        printf("Cannot open input data file for reading\n");
        exit(1);
    }

	// print all requested packets in the pcap file
	if (arg_parser.get<int>("print") > 0 ) {
		PrintPcapPackets(dynamic_cast<pcpp::PcapFileReaderDevice*>(reader), &cout, arg_parser.get<int>("print"));
		return 0;
	};

	// handle packets
    auto loop_start_time = std::chrono::steady_clock::now(); // record loop start time
	PcapPacketsHandler(dynamic_cast<pcpp::PcapFileReaderDevice*>(reader), opts);
    auto loop_end_time = std::chrono::steady_clock::now(); // record loop end time

    // print out use time
	std::chrono::duration<double> time_used = std::chrono::duration_cast<std::chrono::duration<double>>(loop_end_time - loop_start_time);
	std::cout << "Time used: " << time_used.count() << " s" << std::endl;

    // read stats from reader and print them
    pcap_stat status; // create the status object
    reader->getStatistics(status);
    printf("Read %d packets successfully and %d packets could not be read\n", status.ps_recv, status.ps_drop);

	reader->close(); // close the file
	delete reader; // free reader memory

	return 0;
}